from pytest import MonkeyPatch
from fastapi import FastAPI, status, Request
from fastapi.testclient import TestClient
import pytest
import logging

from motor.config.coordinator import DeployMode, CoordinatorConfig
from motor.coordinator.core.instance_manager import InstanceManager
from motor.coordinator.router.base_router import ScheduledResource
from motor.coordinator.router.pd_hybrid_router import PDHybridRouter
from motor.resources.instance import Endpoint, PDRole, Instance, InsStatus, ParallelConfig
from motor.coordinator.scheduler.scheduler import Scheduler
from motor.coordinator.models.request import RequestInfo, ReqState
import motor.coordinator.router.router as router

logger = logging.getLogger(__name__)

app = FastAPI()
@app.post("/v1/chat/completions")
async def handle_completions(request: Request):
    return await router.handle_request(request)

@pytest.fixture
def mock_forward_stream_request(monkeypatch):
    """Mock forward_stream_request 并自动设置和清理"""
    async def mock_impl(self, req_data: dict, resource):
        async def mock_stream():
            responses = [
                b'{"choices": [{"text": "chunk 1"}]}',
                b'{"choices": [{"text": "chunk 2"}]}',
                b'{"choices": [{"text": "chunk 3"}]}'
            ]
            for response in responses:
                yield response
        
        async for chunk in mock_stream():
            yield chunk
    
    # Patch the forward_stream_request function to return an async generator directly
    monkeypatch.setattr(PDHybridRouter, "forward_stream_request", mock_impl)
    yield mock_impl

class TestRouterPDHybrid:
    
    @pytest.fixture
    def client(self):
        return TestClient(app)
    
    def create_mock_instance(self, instance_id, role):
        """Create a proper mock Instance object"""
        mock_instance = Instance(
            job_name=f"test-job-{instance_id}",
            model_name=f"test-model-{instance_id}",
            id=instance_id,
            role=role,
            status=InsStatus.ACTIVE,
            parallel_config=ParallelConfig(dp_size=1, tp_size=1),
            endpoints={}
        )
        return mock_instance
    
    @pytest.fixture
    def setup_pd_hybrid(self, monkeypatch: MonkeyPatch):
        # Create proper instance for PD hybrid flow
        mock_instance_u = self.create_mock_instance(0, PDRole.ROLE_U)
        mock_endpoint_u = Endpoint(id=0, ip="127.0.0.1", port="8000")
        mock_instance_u.endpoints = {"127.0.0.1": {0: mock_endpoint_u}}
        
        # Mock functions
        def mock_is_available(self):
            return True
        
        def mock_get_available_instances(role):
            if role == PDRole.ROLE_U:  # PD hybrid role
                return [mock_instance_u]  # Has PD hybrid instances
            return []
        
        def mock_select_instance_and_endpoint(self, role):
            if role == PDRole.ROLE_U:
                return mock_instance_u, mock_endpoint_u
            return None, None, None
        
        def mock_update_workload(self, instance: Instance, endpoint: Endpoint, req_id: str,
                        workload_action, request_length: int) -> bool:
            return True
        
        def mock_from_string(value):
            return DeployMode.SINGLE_NODE
        
        monkeypatch.setattr(InstanceManager, "is_available", mock_is_available)
        monkeypatch.setattr(InstanceManager, "get_available_instances", mock_get_available_instances)
        monkeypatch.setattr(Scheduler, "select_instance_and_endpoint", mock_select_instance_and_endpoint)
        monkeypatch.setattr(Scheduler, "update_workload", mock_update_workload)
        monkeypatch.setattr(DeployMode, "from_string", mock_from_string)
    
    @pytest.mark.asyncio
    async def test_pd_hybrid_request_forwarding(self, monkeypatch: MonkeyPatch, setup_pd_hybrid, mock_forward_stream_request):
        """Test PD hybrid request forwarding functionality"""
        # Create a RequestInfo
        from fastapi import Request
        import json
        
        # Create a mock scope for the request
        scope = {
            "type": "http",
            "method": "POST",
            "path": "/v1/chat/completions",
            "headers": [],
        }
        
        # Create a request body
        request_body = {
            "model": "test-model",
            "messages": [{"role": "user", "content": "Hello"}],
            "stream": True
        }
        
        # Create a mock request object
        request = Request(scope)
        request._body = json.dumps(request_body).encode()
        
        request_body = await request.body()
        req_len = len(request_body)
        request_json = await request.json()
        
        # Create a RequestInfo
        req_info = RequestInfo(
            req_id="test-id",
            req_data=request_json.copy(),
            req_len=req_len,
            api="v1/chat/completions",
            state=ReqState.ARRIVE
        )
        
        # Test the PD hybrid forwarding function
        router = PDHybridRouter(req_info)
        chunks = []
        
        response = await router.handle_request()
        async for chunk in response.body_iterator:
            chunks.append(chunk)
        
        # Verify we got response chunks
        assert len(chunks) > 0
        # Verify request state was updated
        assert req_info.state == ReqState.DECODE_END
    
    @pytest.mark.asyncio
    async def test_pd_hybrid_request_failure(self, monkeypatch: MonkeyPatch, setup_pd_hybrid):
        """Test handling of PD hybrid request failure"""
        # Create a RequestInfo
        from fastapi import Request
        import json
        
        # Create a mock scope for the request
        scope = {
            "type": "http",
            "method": "POST",
            "path": "/v1/chat/completions",
            "headers": [],
        }
        
        # Create a request body
        request_body = {
            "model": "test-model",
            "messages": [{"role": "user", "content": "Hello"}],
            "stream": True
        }
        
        # Create a mock request object
        request = Request(scope)
        request._body = json.dumps(request_body).encode()
        request_body = await request.body()
        req_len = len(request_body)
        request_json = await request.json()
        
        # Create a RequestInfo
        req_info = RequestInfo(
            req_id="test-id",
            req_data=request_json.copy(),
            req_len=req_len,
            api="v1/chat/completions",
            state=ReqState.ARRIVE
        )
        
        # Mock the stream request function to fail in PDHybridRouter
        async def mock_forward_stream_request(self, req_data, resource):
            # This function should be an async generator that raises an exception
            raise Exception("PD hybrid request failed")
            yield  # This line will never be reached, but makes it a generator
        monkeypatch.setattr(PDHybridRouter, "forward_stream_request", mock_forward_stream_request)
        
        # Test the PD hybrid forwarding function with failure
        router = PDHybridRouter(req_info)
        with pytest.raises(Exception) as exc_info:
            # Create an async generator and consume it
            stream_resp = await router.handle_request()
            async for chunk in stream_resp.body_iterator:
                pass
        
        assert "PD hybrid request failed" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_successful_request_with_pd_hybrid(self, client, monkeypatch: MonkeyPatch, setup_pd_hybrid):
        """测试场景:PD混合模式请求成功
        预期现象:
        1)检查请求状态为DecodeEnd
        2)返回正常响应
        """
        # Mock the HTTP forwarding function to return a successful response
        async def mock_forward_stream_request(self, req_data: dict, resource: ScheduledResource):
            # Yield a simple response
            yield b'{"choices": [{"delta": {"content": "Hello"}}]}'
        
        monkeypatch.setattr(PDHybridRouter, "forward_stream_request", mock_forward_stream_request)
        
        response = client.post("/v1/chat/completions", json={
            "model": "test-model", 
            "messages": [{"role": "user", "content": "Hello"}],
            "stream": True
        })
        
        # Should get a 200 success status
        assert response.status_code == status.HTTP_200_OK
        # Should be a streaming response
        assert response.headers.get("content-type") == "application/json"
